
from space import *
from numpy import array, arange, prod

class IndicesMappingSpace(MappingSpace):
    """This Space represents a function that can convert a list of integer into an integer
        representing the indice of elements in a n-dimension matrix
        
        This mapping is used to pass from an n-dimension array to a 1 dimension array
    """
    def __init__(self, n):
        MappingSpace.__init__(self, natural**n, natural)

    class Element(MappingSpace.Element):
        def __init__(self, space, dims):
            self.dims = dims
            # We store all the index into an n dimension array
            # (may not be efficient for big arrays)
            self.index = arange(self.modulo(0))
            self.index.resize(dims)
        
            inverse = space.inverse()(self.inv_map)
            MappingSpace.Element.__init__(self, space, self.map, inverse)
            
        def modulo(self, n):
            """Return the length between 2 consecutif values along a given dimention"""
            return prod(self.dims[n:])
        
        def map(self, indices):
            """Return the index of the indices"""
            # We use the precomputed index array
            return self.index[tuple(indices)]
        
        def inv_map(self, index):
            ret = []
            for d in self.dims:
                ret.append(int(index / d))
                index -= ret[-1]
            return ret
            

class ProdMappingSpace(MappingSpace):
    """
    Combine several mapping space into a single ProdSpace Mapping
    
    This create a mapping space : (S1 * S2 * ...) -> (D1 * D2 * ...)
    from several mapping spaces S1 -> D1, S2 -> D2, ...
    """
    def __init__(self, mappings):
        self.mappings = mappings
        src = ProdSpace([m.src for m in mappings])
        dst = ProdSpace([m.dst for m in mappings])
        MappingSpace.__init__(self, src, dst)    

    class Element(MappingSpace.Element):
        def __init__(self, space, mappings):
            self.mappings = mappings
            MappingSpace.Element.__init__(self, space, self.map)
        
        def map(self, coords):
            return [m(v).coords() for (m,v) in zip(self.mappings, coords)]
            
            
class ToIndexMappingSpace(MappingSpace):
    """Mapping space : src -> natural, that works in a linear way
    
        This mapping is used for the ImageSpace, where we need to map a coordinate into an index
    """
    def __init__(self, src):
        MappingSpace.__init__(self, src, natural)
    class Element(MappingSpace.Element):
        def __init__(self, space, n, min, max):
            """Create a linear mapping to a natural
            
            @param n     the number of index
            @param min   the min value (at index 0)
            @param max   the max value (at index N-1)
            """
            self.n, self.min, self.max = n, min, max
            inverse = space.inverse()(self.inv_map)
            MappingSpace.Element.__init__(self, space, self.map, inverse)
        def map(self, v):
            return int((v - self.min) / (self.max - self.min) * self.n)
        def inv_map(self, i):
            return self.min + i * (self.max - self.min) / self.n


class ImageMapping(DiscreteMappingSpace):
    """Specialized mapping to describe n dimensional images"""
    def __init__(self, space, shape, name = None):
        """
        @param space  the mapping space describing the image
                      space.src must be a ProductSpace
        @param shape   a list of (n, min, max) values, defining the dimension of the image
        """
        self.shape = shape
        n = prod([d[0] for d in shape])
        
        if isinstance(space.src, ProdSpace):
            src_spaces = space.src.spaces
            to_indexes = [ToIndexMappingSpace(s)(d[0], d[1], d[2]) for (s,d) in zip(src_spaces, shape)]
            
            v_to_is = ProdMappingSpace([e.space for e in to_indexes])(to_indexes)
            is_to_i = IndicesMappingSpace(len(shape))([d[0] for d in shape])
            v_to_i = v_to_is * is_to_i
        else:   # one dimension image
            d = shape[0]
            v_to_i = ToIndexMappingSpace(space.src)(d[0], d[1], d[2])
        
        DiscreteMappingSpace.__init__(self, space, n, v_to_i, lambda x:1, name = name)
    



if __name__ == '__main__':
    print '== indices mapping =='
    indices_s = IndicesMappingSpace(2)
    indices = indices_s([2,2])
    print '[1,0] ->', indices([1,0]).coords()
    print '2 ->', indices.inverse(2).coords()
    
    print '== ProdMapping =='
    m1 = MappingSpace(real, real)(lambda x: 2 *x)
    m2 = MappingSpace(real, real)(lambda x: x ** 2)
    m = ProdMappingSpace([m1.space, m2.space])([m1, m2])
    print m([1,1])
    
    print '== ToIndexMappingSpace =='
    to_index = ToIndexMappingSpace(real)(10, 0., 5.)
    print '0.5 ->', to_index(0.5)
    print '1 ->', to_index.inverse(1)
    
    print '== 1D image =='
    image_data = array([1,2,3,4])
    pos_s = real
    func_space = MappingSpace(pos_s, real)
    shape = ((4,0,1),)
    image_s = ImageMapping(func_space, shape)
    image = image_s(image_data)  # TODO: avoid the flat
    print "image(0.9) = ", image(0.9)
    
    print '== 2D image =='
    image_data = array([[1,2],[3,4]])
    pos_s = ProdSpace([real, real])
    func_space = MappingSpace(pos_s, real)
    shape = ((2,0,1),(2,0,1))  # both dimensions from 0 to 1, with 2 values per dimension
    image_s = ImageMapping(func_space, shape)
    image = image_s(image_data.flat)  # TODO: avoid the flat
    print "image([0,0]) = ", image([0,0])
